"""
WebSocket连接管理器

统一的WebSocket连接管理，支持：
1. 连接数限制和心跳检测
2. 主题订阅和消息广播
3. 连接状态监控和清理
4. 用户身份管理
"""

import asyncio
import time
from contextlib import asynccontextmanager
from typing import Set, Dict, Any, Optional

from fastapi import WebSocket, WebSocketDisconnect

from common.config import settings
from common.logger import create_logger

logger = create_logger(
	__name__,
	log_level=settings.log_level,
	log_file=settings.log_file,
	enable_file_logging=settings.enable_file_logging
)

# WebSocket连接限制配置
MAX_WEBSOCKET_CONNECTIONS = 500  # 最大连接数
HEARTBEAT_INTERVAL = 20  # 心跳间隔（秒）
CONNECTION_TIMEOUT = 300  # 连接超时时间（秒）


class WebSocketConnectionManager:
	"""WebSocket连接管理器 - 统一版本"""

	def __init__(self, max_connections: int = MAX_WEBSOCKET_CONNECTIONS):
		self.max_connections = max_connections
		self.active_connections: Set[WebSocket] = set()
		self.user_connections: Dict[str, WebSocket] = {}
		self.connection_metadata: Dict[WebSocket, Dict[str, Any]] = {}
		self._cleanup_task: Optional[asyncio.Task] = None
		self._lock = asyncio.Lock()

	async def connect(self, websocket: WebSocket, user_code: str = None) -> bool:
		"""
		建立WebSocket连接

		Args:
			websocket: WebSocket连接对象
			user_code: 用户标识（可选）

		Returns:
			bool: 是否连接成功
		"""
		async with self._lock:
			# 检查连接数限制
			if len(self.active_connections) >= self.max_connections:
				logger.warning(f"WebSocket连接数已达上限 {self.max_connections}，拒绝新连接")
				try:
					await websocket.close(code=1008, reason="连接数已达上限")
				except:
					pass
				return False

			try:
				await websocket.accept()
				self.active_connections.add(websocket)

				# 记录连接元数据
				current_time = time.time()
				self.connection_metadata[websocket] = {
					"connected_at": current_time,
					"last_ping": current_time,
					"client_ip": websocket.client.host if websocket.client else "unknown",
					"user_code": user_code,
					"subscriptions": set()  # 订阅的主题
				}

				if user_code:
					# 如果用户已有连接，先断开旧连接
					if user_code in self.user_connections:
						old_websocket = self.user_connections[user_code]
						await self._force_disconnect(old_websocket)
					self.user_connections[user_code] = websocket

				logger.info(f"WebSocket连接建立 (IP: {self.connection_metadata[websocket]['client_ip']}, "
							f"用户: {user_code or 'anonymous'})，当前连接数: {len(self.active_connections)}")

				# 启动清理任务
				if self._cleanup_task is None or self._cleanup_task.done():
					self._cleanup_task = asyncio.create_task(self._periodic_cleanup())

				return True

			except Exception as e:
				logger.error(f"WebSocket连接建立失败: {e}")
				return False

	async def disconnect(self, websocket: WebSocket, user_code: str = None):
		"""
		断开WebSocket连接

		Args:
			websocket: WebSocket连接对象
			user_code: 用户标识（可选）
		"""
		async with self._lock:
			self._remove_connection(websocket, user_code)

	async def send_personal_message(self, message: dict, user_code: str) -> bool:
		"""
		发送个人消息

		Args:
			message: 消息内容
			user_code: 用户标识

		Returns:
			bool: 是否发送成功
		"""
		if user_code not in self.user_connections:
			return False

		websocket = self.user_connections[user_code]
		return await self._send_safe(websocket, message)

	async def broadcast(self, message: dict, topic: str = None, exclude_user: str = None) -> int:
		"""
		广播消息给所有连接

		Args:
			message: 消息内容
			topic: 主题过滤（可选）
			exclude_user: 排除的用户（可选）

		Returns:
			int: 成功发送的连接数
		"""
		if not self.active_connections:
			return 0

		target_connections = []
		for websocket in list(self.active_connections):  # 使用副本避免并发修改
			metadata = self.connection_metadata.get(websocket, {})

			# 排除指定用户
			if exclude_user and metadata.get("user_code") == exclude_user:
				continue

			# 如果指定了主题，检查订阅状态
			if topic is None or topic in metadata.get("subscriptions", set()):
				target_connections.append(websocket)

		if not target_connections:
			logger.info(f"广播消息但无订阅连接 (主题: {topic or '全部'})")
			return 0

		logger.info(f"广播WebSocket消息到 {len(target_connections)} 个连接 (主题: {topic or '全部'})")

		# 并发发送消息
		tasks = [self._send_safe(conn, message) for conn in target_connections]
		results = await asyncio.gather(*tasks, return_exceptions=True)

		success_count = sum(1 for r in results if r is True)
		logger.info(f"WebSocket消息发送完成: 成功 {success_count}/{len(tasks)}")

		return success_count

	async def _send_safe(self, websocket: WebSocket, message: dict) -> bool:
		"""
		安全发送WebSocket消息

		Args:
			websocket: WebSocket连接对象
			message: 消息内容

		Returns:
			bool: 是否发送成功
		"""
		try:
			await websocket.send_json(message)
			# 更新最后活跃时间
			if websocket in self.connection_metadata:
				self.connection_metadata[websocket]["last_ping"] = time.time()
			return True
		except Exception as e:
			logger.warning(f"WebSocket消息发送失败: {e}")
			# 移除失效连接
			await self._force_disconnect(websocket)
			return False

	def _remove_connection(self, websocket: WebSocket, user_code: str = None):
		"""
		移除连接（内部方法）

		Args:
			websocket: WebSocket连接对象
			user_code: 用户标识（可选）
		"""
		self.active_connections.discard(websocket)
		metadata = self.connection_metadata.pop(websocket, {})

		# 从用户连接映射中移除
		actual_user_code = user_code or metadata.get("user_code")
		if actual_user_code and actual_user_code in self.user_connections:
			if self.user_connections[actual_user_code] == websocket:
				del self.user_connections[actual_user_code]

		logger.info(f"WebSocket连接断开 (IP: {metadata.get('client_ip', 'unknown')}, "
					f"用户: {actual_user_code or 'anonymous'})，当前连接数: {len(self.active_connections)}")

	async def _force_disconnect(self, websocket: WebSocket):
		"""
		强制断开连接

		Args:
			websocket: WebSocket连接对象
		"""
		try:
			await websocket.close(code=1000, reason="连接被替换或清理")
		except:
			pass
		self._remove_connection(websocket)

	async def _periodic_cleanup(self):
		"""定期清理僵尸连接"""
		while True:
			try:
				await asyncio.sleep(60)  # 每分钟检查一次
				await self._cleanup_stale_connections()
			except asyncio.CancelledError:
				logger.info("WebSocket清理任务被取消")
				break
			except Exception as e:
				logger.error(f"清理连接任务异常: {e}")

	async def _cleanup_stale_connections(self):
		"""清理僵尸连接"""
		current_time = time.time()
		stale_connections = []

		async with self._lock:
			for websocket, metadata in list(self.connection_metadata.items()):
				# 超过CONNECTION_TIMEOUT秒无活动的连接视为僵尸连接
				if current_time - metadata.get("last_ping", 0) > CONNECTION_TIMEOUT:
					stale_connections.append(websocket)

		# 清理僵尸连接
		for websocket in stale_connections:
			await self._force_disconnect(websocket)

		if stale_connections:
			logger.info(f"清理了 {len(stale_connections)} 个僵尸WebSocket连接")

	def subscribe_topic(self, websocket: WebSocket, topic: str):
		"""
		订阅主题

		Args:
			websocket: WebSocket连接对象
			topic: 主题名称
		"""
		if websocket in self.connection_metadata:
			self.connection_metadata[websocket]["subscriptions"].add(topic)
			logger.debug(f"WebSocket连接订阅主题: {topic}")

	def unsubscribe_topic(self, websocket: WebSocket, topic: str):
		"""
		取消订阅主题

		Args:
			websocket: WebSocket连接对象
			topic: 主题名称
		"""
		if websocket in self.connection_metadata:
			self.connection_metadata[websocket]["subscriptions"].discard(topic)
			logger.debug(f"WebSocket连接取消订阅主题: {topic}")

	def get_connection_stats(self) -> Dict[str, Any]:
		"""
		获取连接统计信息

		Returns:
			Dict: 连接统计信息
		"""
		current_time = time.time()

		connection_stats = []
		for websocket, metadata in self.connection_metadata.items():
			connection_stats.append({
				"client_ip": metadata.get("client_ip", "unknown"),
				"user_code": metadata.get("user_code"),
				"connected_duration": int(current_time - metadata.get("connected_at", current_time)),
				"last_activity": int(current_time - metadata.get("last_ping", current_time)),
				"subscriptions": list(metadata.get("subscriptions", set()))
			})

		return {
			"total_connections": len(self.active_connections),
			"max_connections": self.max_connections,
			"user_connections": len(self.user_connections),
			"connections": connection_stats
		}

	async def close_all_connections(self, reason: str = "服务器关闭"):
		"""
		关闭所有连接

		Args:
			reason: 关闭原因
		"""
		logger.info(f"开始关闭所有WebSocket连接，原因: {reason}")

		# 发送关闭通知
		if self.active_connections:
			close_message = {
				"type": "server_shutdown",
				"message": reason,
				"timestamp": time.time()
			}
			await self.broadcast(close_message)

			# 等待客户端处理消息
			await asyncio.sleep(1)

			# 强制关闭所有连接
			for websocket in list(self.active_connections):
				await self._force_disconnect(websocket)

		# 取消清理任务
		if self._cleanup_task and not self._cleanup_task.done():
			self._cleanup_task.cancel()
			try:
				await self._cleanup_task
			except asyncio.CancelledError:
				pass

		logger.info("所有WebSocket连接已关闭")


# 全局WebSocket管理器实例
websocket_manager = WebSocketConnectionManager()


def get_websocket_manager() -> WebSocketConnectionManager:
	"""获取WebSocket管理器实例"""
	return websocket_manager


@asynccontextmanager
async def websocket_connection(websocket: WebSocket, user_code: str = None):
	"""
	WebSocket连接上下文管理器

	Args:
		websocket: WebSocket连接对象
		user_code: 用户标识（可选）
	"""
	manager = get_websocket_manager()

	if not await manager.connect(websocket, user_code):
		return

	try:
		yield manager
	except WebSocketDisconnect:
		logger.info("WebSocket连接主动断开")
	except Exception as e:
		logger.error(f"WebSocket连接异常: {e}")
	finally:
		await manager.disconnect(websocket, user_code)
